#pragma once

#include "viml/value.h"
#include "viml/ast.h"
#include "viml/instruction.h"
#include <memory>
#include <vector>
#include <unordered_set>
#include <unordered_map>
#include <cmath>

namespace viml {

// 前向声明
class BinaryExpression;
class UnaryExpression;
class ASTVisitor;

/**
 * 常量折叠访问器
 * 遍历AST并执行常量表达式求值
 */
class ConstantFolder : public ASTVisitor {
private:
    std::unique_ptr<ASTNode> result_;  // 改用unique_ptr支持clone
    bool changed_ = false;
    size_t folds_ = 0;

public:
    ConstantFolder() = default;

    std::unique_ptr<ASTNode> get_result() { return std::move(result_); }  // 返回unique_ptr
    bool has_changed() const { return changed_; }
    size_t get_fold_count() const { return folds_; }

    // 重置状态
    void reset() {
        result_ = nullptr;
        changed_ = false;
        folds_ = 0;
    }

    // 访问者方法
    void visit(NumberLiteral& node) override;
    void visit(StringLiteral& node) override;
    void visit(Identifier& node) override;
    void visit(SpecialConstantNode& node) override;
    void visit(BinaryExpression& node) override;
    void visit(UnaryExpression& node) override;
    void visit(CallExpression& node) override;
    void visit(ListLiteral& node) override;
    void visit(DictLiteral& node) override;
    void visit(IndexExpression& node) override;
    void visit(LetStatement& node) override;
    void visit(EchoStatement& node) override;
    void visit(IfStatement& node) override;
    void visit(WhileStatement& node) override;
    void visit(ForStatement& node) override;
    void visit(FunctionStatement& node) override;
    void visit(ReturnStatement& node) override;
    void visit(CallStatement& node) override;
    void visit(Program& node) override;

private:
    // 常量求值辅助方法
public:  // 为了SimpleOptimizer访问，设为public
    std::shared_ptr<Expression> evaluate_binary_op(BinaryExpression& expr);
    std::shared_ptr<Expression> evaluate_unary_op(UnaryExpression& expr);
    bool is_constant(const std::shared_ptr<Expression>& expr);
    Value evaluate_expression(const std::shared_ptr<Expression>& expr);
    Value evaluate_expression(const std::unique_ptr<Expression>& expr);
};

/**
 * 死代码消除器
 * 识别并移除不可达代码
 */
class DeadCodeEliminator {
private:
    std::unordered_set<size_t> reachable_statements_;
    bool changed_ = false;
    size_t eliminated_count_ = 0;

public:
    DeadCodeEliminator() = default;

    bool has_changed() const { return changed_; }
    size_t get_eliminated_count() const { return eliminated_count_; }

    // 分析并消除死代码
    std::shared_ptr<Program> eliminate_dead_code(const std::shared_ptr<Program>& program);

private:
    void analyze_reachability(const std::shared_ptr<Program>& program);
    void mark_reachable(const std::shared_ptr<ASTNode>& node);
};

/**
 * 简化的优化器
 * 实现基本的常量折叠、死代码消除和寄存器重用优化
 */
class SimpleOptimizer {
private:
    bool enabled_ = true;

    // 优化统计
    struct OptimizationCounters {
        size_t constant_folds = 0;
        size_t dead_code_eliminated = 0;
        size_t register_reuses = 0;
        size_t optimizations_total = 0;
    } counters_;

public:
    SimpleOptimizer() = default;

    /**
     * 优化AST
     */
    std::shared_ptr<Program> optimize_ast(const std::shared_ptr<Program>& program) {
        if (!enabled_) {
            return program;
        }

        std::shared_ptr<Program> result = program;

        // 1. 常量折叠优化
        result = optimize_constants_in_ast(result);

        // 2. 死代码消除
        result = eliminate_dead_code(result);

        // 3. 寄存器重用优化（在字节码级别进行）
        // 这里先统计，实际重用在字节码优化阶段

        counters_.optimizations_total++;
        return result;
    }

    /**
     * 优化字节码
     */
    std::shared_ptr<Bytecode> optimize_bytecode(const std::shared_ptr<Bytecode>& bytecode) {
        if (!enabled_) {
            return bytecode;
        }

        // 这里可以添加字节码级别的优化
        // 1. 寄存器重用优化
        // 2. 指令简化等
        optimize_registers_in_bytecode(bytecode);

        return bytecode;
    }

    /**
     * 启用/禁用优化
     */
    void set_enabled(bool enabled) {
        enabled_ = enabled;
    }

    bool is_enabled() const {
        return enabled_;
    }

    /**
     * 获取优化统计信息
     */
    struct OptimizationSummary {
        size_t constant_folds = 0;
        size_t dead_code_eliminated = 0;
        size_t register_reuses = 0;
        size_t optimizations_total = 0;
        double improvement_factor = 0.0;
    };

    OptimizationSummary get_optimization_summary() const {
        OptimizationSummary summary;
        summary.constant_folds = counters_.constant_folds;
        summary.dead_code_eliminated = counters_.dead_code_eliminated;
        summary.register_reuses = counters_.register_reuses;
        summary.optimizations_total = counters_.optimizations_total;

        if (counters_.optimizations_total > 0) {
            summary.improvement_factor = 100.0;  // 简化的改进度量
        }

        return summary;
    }

    /**
     * 重置优化统计
     */
    void reset_counters() {
        counters_ = OptimizationCounters{};
    }

    // 公开优化方法，供 Compiler 调用
    std::shared_ptr<Program> optimize_constants_in_ast(const std::shared_ptr<Program>& program);
    std::shared_ptr<Program> eliminate_dead_code(const std::shared_ptr<Program>& program);
    void optimize_registers_in_bytecode(const std::shared_ptr<Bytecode>& bytecode);

private:

    // 优化辅助方法
    bool is_constant_expression(const std::shared_ptr<Expression>& expr);
    std::shared_ptr<Expression> fold_constants(const std::shared_ptr<Expression>& expr);
    Value evaluate_constant_expression(const std::shared_ptr<Expression>& expr);
};

} // namespace viml